package com.inspur.dgov.vlbod.util;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;

import java.util.Arrays;
import java.util.Collections;
import java.util.List;

/**
 * LimitUtils
 * 限流
 * package com.inspur.dgov.vlbod.util
 *
 * @author xie_yh
 * @date 2023/6/5 21:08
 */
@Component
public class RateLimiterUtils {

    @Autowired
    private RedisTemplate redisTemplate;

    public boolean tryAcquire(String prefix,String key,String tag,int limitCount,int limitPeriod){
        List<Object> keys = Collections.singletonList(prefix + "-" + key + "-" + tag);

        String luaScript = buildLuaScript();
        RedisScript<Long> redisScript = new DefaultRedisScript<>(luaScript, Long.class);
        Long count = (Long) redisTemplate.execute(redisScript, keys, limitCount, limitPeriod);
        if (null != count && count.intValue() <= limitCount) {
            return true;
        } else {
            return false;
        }
    }

    private String buildLuaScript() {
        return "local key = KEYS[1]\n" +
                "local count = tonumber(ARGV[1])\n" +
                "local time = tonumber(ARGV[2])\n" +
                "local current = redis.call('get', key);\n" +
                "if current and tonumber(current) > count then\n" +
                "    return tonumber(current);\n" +
                "end\n" +
                "current = redis.call('incr', key)\n" +
                "if tonumber(current) == 1 then\n" +
                "    redis.call('expire', key, time)\n" +
                "end\n" +
                "return tonumber(current);";
    }

}
